package com.shop.cloud.common.data.tenant;

import cn.hutool.core.util.StrUtil;
import com.alibaba.ttl.TransmittableThreadLocal;
import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.creator.DataSourceProperty;
import com.baomidou.dynamic.datasource.creator.DefaultDataSourceCreator;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.shop.cloud.common.core.constant.CommonConstants;
import com.shop.cloud.common.core.exception.CheckedException;
import com.shop.cloud.common.core.util.SpringContextHolder;
import com.shop.cloud.common.data.shardingdatabase.config.ShardingDataBaseProperties;
import lombok.Synchronized;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import javax.sql.DataSource;
import java.sql.DriverManager;

/**
 * @author
 * 租户工具类
 */
@Slf4j
@UtilityClass
public class TenantContextHolder {

	private final ThreadLocal<String> THREAD_LOCAL_TENANT = new TransmittableThreadLocal<>();
	private DataSource dataSource = SpringContextHolder.getBean(DataSource.class);
	private DefaultDataSourceCreator dataSourceCreator = SpringContextHolder.getBean(DefaultDataSourceCreator.class);
	private ShardingDataBaseProperties shardingDataBaseProperties = SpringContextHolder.getBean(ShardingDataBaseProperties.class);

	/**
	 * TTL 设置租户ID
	 *
	 * @param tenantId
	 */
	public void setTenantId(String tenantId) {
		THREAD_LOCAL_TENANT.set(tenantId);
		//分库实现，根据租户ID，切换当前租户的数据源
		if(shardingDataBaseProperties != null && shardingDataBaseProperties.getEnabled() != null
				&& shardingDataBaseProperties.getEnabled() && StrUtil.isNotBlank(tenantId) && !"null".equals(tenantId)){
			DynamicRoutingDataSource ds = (DynamicRoutingDataSource) dataSource;
			//拼接当前租户的数据源名
			String dsName;
			if(CommonConstants.SYSTEM_TENANT_ID.equals(tenantId)){
				dsName = shardingDataBaseProperties.getDbname();
			}else{
				dsName = StrUtil.format("{}_{}", shardingDataBaseProperties.getDbname(), tenantId);
			}
			String dbName = dsName;
			//如数据源中没有该数据源，则添加
			if(!ds.getDataSources().containsKey(dsName)) {
				addDataSources(ds, dataSourceCreator, dsName, dbName, shardingDataBaseProperties);
			}
			//切换到当前租户数据源
			DynamicDataSourceContextHolder.push(dsName);
		}
	}

	/**
	 * 获取TTL中的租户ID
	 *
	 * @return
	 */
	public String getTenantId() {
		return THREAD_LOCAL_TENANT.get();
	}

	public void clear() {
		THREAD_LOCAL_TENANT.remove();
	}

	/**
	 * 添加数据源，如数据源中没有该数据源，则添加
	 * @param dataSourceCreator
	 * @param dsName
	 * @param dbName
	 * @param shardingDataBaseProperties
	 * @return
	 */
	@Synchronized
	public void addDataSources(DynamicRoutingDataSource ds, DefaultDataSourceCreator dataSourceCreator,
							   String dsName, String dbName, ShardingDataBaseProperties shardingDataBaseProperties){
		if(!ds.getDataSources().containsKey(dsName)) {
			//获取主库账号密码url
			String pwd = shardingDataBaseProperties.getDbpassword();
			String username = shardingDataBaseProperties.getDbusername();
			String url = StrUtil.format(shardingDataBaseProperties.getUrl(), dbName);
			//检查数据源是否可用
			try {
				DriverManager.getConnection(url, username, pwd);
			} catch (Exception e) {
				log.error("数据源不可用："+e.getMessage());
				throw new CheckedException("数据源不可用："+e.getMessage());
			}
			DataSourceProperty dataSourceProperty = new DataSourceProperty();
			dataSourceProperty.setPoolName(dsName);
			dataSourceProperty.setPassword(pwd);
			dataSourceProperty.setUsername(username);
			dataSourceProperty.setUrl(url);
			//添加数据源
			DataSource dataSource = dataSourceCreator.createDataSource(dataSourceProperty);
			ds.addDataSource(dataSourceProperty.getPoolName(), dataSource);
		}
	}
}
